# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: taekwon

import datetime
import numpy as np
import pandas as pd
import json
import fire

from coin.base.query_util import (query_klines, query_pta)


def get_market_share(start_date, end_date, business_unit, breakdown):
  print(f"{start_date} ~ {end_date}")
  curr_date = datetime.datetime.utcnow().date()
  end_date = datetime.datetime.strptime(end_date, '%Y%m%d').date() \
      if end_date != '' else curr_date - datetime.timedelta(2)
  start_date = datetime.datetime.strptime(start_date, '%Y%m%d').date() \
      if start_date != '' else end_date
  assert start_date <= end_date, (start_date, end_date)
  pta_df = query_pta(start_date, end_date, business_units=[business_unit], pta_type='SYMBOL')
  pta_df['fill_pq_in_usd'] = pta_df['fill_pq_in_usd'].astype('float')
  dft_group = ['market_type', 'exchange', 'symbol']
  valid_col = ['trader', 'strategy_name', 'strategy_group', 'fill_pq_in_usd']
  pta_df = pta_df[dft_group + valid_col]
  turnover_df = pta_df.groupby(dft_group + breakdown).sum().reset_index()
  
  klines = query_klines(start_date, end_date)
  klines['external_symbol'] = klines['kline_dict'].apply(lambda x: json.loads(x).get('symbol'))
  klines['turnover_in_usd'] = klines['kline_dict'].apply(lambda x: (json.loads(x).get('klines', [{}]))[0].get('turnover_in_usd'))
  rklines = klines[['market_type','exchange','external_symbol','turnover_in_usd']]
  rklines = rklines.groupby(['market_type', 'exchange', 'external_symbol']).sum().reset_index()
  rklines.rename(columns={'external_symbol': 'symbol'}, inplace=True)
  res_df = pd.merge(turnover_df, rklines, on=dft_group, how='left')
  res_df['market_share'] = res_df['fill_pq_in_usd'] * 100.0 / res_df['turnover_in_usd']
  return res_df.sort_values(by='market_share', ascending=False)


def filter_result(res_df, filter_dict):
  for ft_name, ft_val in filter_dict.items():
    if ft_val is not None:
      res_df = res_df[res_df[ft_name] == ft_val]
  return res_df


def print_market_share(trading_date,
                       market_type=None, exchange=None, symbol=None,
                       breakdown=None, to_csv=False, business_unit='Coin'):
  trading_date_list = str(trading_date).split("-")
  if len(trading_date_list) == 1:
    start_date = trading_date_list[0]
    end_date = trading_date_list[0]
  else:
    start_date, end_date = trading_date_list
  if breakdown is not None:
    breakdown = breakdown.split(',')
  else:
    breakdown = list()
  res_df = get_market_share(start_date, end_date, business_unit, breakdown)
  res_df.rename(columns={'market_share': 'market_share(%)'}, inplace=True)

  filter_dict = dict()
  filter_dict['market_type'] = market_type
  filter_dict['exchange'] = exchange
  filter_dict['symbol'] = symbol
  res_df = filter_result(res_df, filter_dict)
  print(res_df[:50])
  if to_csv:
    filename = f'{start_date}_{end_date}_market_share.csv'
    print(f"to csv.. See workspace/coin/{filename}")
    res_df.to_csv(filename, index=False)
    print("Done!")


if __name__ == '__main__':
  fire.Fire()

"""
Usage
./pyrunner python/experimental/taekwon/market_share.py print_market_share --trading_date=20220617-20220619
./pyrunner python/experimental/taekwon/market_share.py print_market_share --trading_date=20220617 --breakdown=strategy_group --to_csv
./pyrunner python/experimental/taekwon/market_share.py print_market_share --trading_date=20220617 --market_type=Futures --exchange=Binance --breakdown=trader --to_csv
"""
